from dataclasses_json import config

from src.module.Model import Model
from langgraph.graph import StateGraph, START, END, MessagesState
from langgraph.checkpoint.memory import MemorySaver

class LangGraphHumanLoop:
    @staticmethod
    def node1(args):
        val = '----> node 1 '
        print(val)
        return { 'messages': val }

    @staticmethod
    def node2(args):
        val = '----> node 2 '
        print(val)
        return { 'messages': val }

    @staticmethod
    def node3(args):
        val = '----> node 3'
        print(val)
        return { 'messages': val }

    def start(self):
        builder = StateGraph(MessagesState)
        builder.add_node('step1', self.node1)
        builder.add_node('step2', self.node2)
        builder.add_node('step3', self.node3)

        builder.add_edge(START, 'step1')
        builder.add_edge('step1', 'step2')
        builder.add_edge('step2', 'step3')
        builder.add_edge('step3', END)

        memory = MemorySaver()
        graph = builder.compile(checkpointer=memory, interrupt_before=['step2'])
        graph.invoke({ "messages": '执行...' }, config={'configurable': {'thread_id': 123}})